package com.jd.jobkeeper.server.transport;

import com.jd.jobkeeper.protocals.Node;
import com.jd.jobkeeper.protocals.Request;
import com.jd.jobkeeper.protocals.ResponseStatus;
import com.jd.jobkeeper.protocals.SubscribeResponse;
import com.jd.jobkeeper.server.data.entity.Namespace;
import com.jd.jobkeeper.server.service.NamespaceService;
import com.jd.jobkeeper.server.service.ServerInfoService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;

import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.stream.Collectors;

@Component
public class JobServerTransportImpl implements JobServerTransport {

    final Logger log = LoggerFactory.getLogger(JobServerTransportImpl.class);
    protected final TcpServer server;

    /**
     * namespace -> nodeId -> nodeContext
     */
    protected final ConcurrentMap<String, ConcurrentNavigableMap<String, NodeContext>> nodeInfos = new ConcurrentHashMap<>();

    protected final Scheduler scheduler = Schedulers.parallel();

    final NamespaceService namespaceService;
    final ServerInfoService serverInfoService;

    public JobServerTransportImpl(TcpServer server,
                                  NamespaceService namespaceService,
                                  ServerInfoService serverInfoService) {
        this.server = server;
        this.namespaceService = namespaceService;
        this.serverInfoService = serverInfoService;
        this.server.registerHandler(this, JobServerTransport.class);
    }

    @Override
    public NodeContext selectNode(String namespace) {
        return nodeInfos.get(namespace).firstEntry().getValue();
    }

    @Override
    public NodeContext selectNode(String namespace, String nodeId) {
        ConcurrentNavigableMap<String, NodeContext> concurrentNavigableMap = nodeInfos.get(namespace);
        if (nodeId == null || nodeId.isEmpty()) {
            return concurrentNavigableMap.firstEntry().getValue();
        } else {
            return concurrentNavigableMap.get(nodeId);
        }
    }

    @Override
    public NodeContext selectNode(String namespace, int index) {
        return null;
    }

    @Override
    public Collection<String> getAllNamespace() {
        return nodeInfos.entrySet().stream()
                .filter(e -> !e.getValue().isEmpty())
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());
    }

    @Override
    public Mono<SubscribeResponse> subscribe(Request request, ServerTransport transport) {
        Sinks.One<SubscribeResponse> one = Sinks.one();
        scheduler.schedule(() -> {
            Namespace namespace = namespaceService.discoveryLeader(request.getNamespace(), transport.getSetupRequest().getAppName());
            if (serverInfoService.isSelf(namespace.getServerId())) {
                String nsCode = request.getNamespace();
                putTransport(nsCode, transport);
                transport.onClose().doFinally(s -> {
                    removeTransport(nsCode, transport);
                }).subscribe();
                one.tryEmitValue(buildResponse(namespace, ResponseStatus.OK));
            } else {
                one.tryEmitValue(buildResponse(namespace, ResponseStatus.ERROR));
            }
        });
        return one.asMono();
    }

    @Override
    public void closeTransport(String namespace) {
        ConcurrentNavigableMap<String, NodeContext> c = nodeInfos.get(namespace);
        if (c != null) {
            c.forEach((nodeId, t) -> {
                t.getServerTransport().dispose();
            });
        }
    }

    private void putTransport(String namespace, ServerTransport serverTransport) {
        NodeContext nodeContext = new NodeContext(serverTransport);
        nodeInfos.computeIfAbsent(namespace, (k) -> new ConcurrentSkipListMap<>())
                .put(nodeContext.nodeId(), nodeContext);
    }

    private void removeTransport(String namespace, ServerTransport serverTransport) {
        log.info("The client is close ! client-> {}", namespace);
        nodeInfos.get(namespace).remove(serverTransport.getSetupRequest().getId());
    }

    private SubscribeResponse buildResponse(Namespace namespace, ResponseStatus status) {
        Node node = Node.newBuilder()
                .setId(namespace.getServerId())
                .setHost(namespace.getHost())
                .setPort(namespace.getPort())
                .build();
        return SubscribeResponse.newBuilder()
                .setServer(node)
                .setStatus(status)
                .build();
    }


}
